#pragma once

#include "imagehlp.h"
#pragma comment(lib, "imagehlp.lib")

static const BYTE cPushOpCode = 0x68;   // The PUSH opcode on x86 platforms



__inline void ReplaceIATEntryInOneMod(PCSTR pszCalleeModName, 
	PROC pfnCurrent, PROC pfnNew, HMODULE hmodCaller) {

	   // Get the address of the module's import section
	   ULONG ulSize;
	   PIMAGE_IMPORT_DESCRIPTOR pImportDesc = (PIMAGE_IMPORT_DESCRIPTOR)
		   ImageDirectoryEntryToData(hmodCaller, TRUE, 
		   IMAGE_DIRECTORY_ENTRY_IMPORT, &ulSize);

	   if (pImportDesc == NULL)
		   return;  // This module has no import section


	   // Find the import descriptor containing references to callee's functions
	   for (; pImportDesc->Name; pImportDesc++) {
		   PSTR pszModName = (PSTR) ((PBYTE) hmodCaller + pImportDesc->Name);
		   if (lstrcmpiA(pszModName, pszCalleeModName) == 0) 
			   break;   // Found
	   }

	   if (pImportDesc->Name == 0)
		   return;  // This module doesn't import any functions from this callee

	   // Get caller's import address table (IAT) for the callee's functions
	   PIMAGE_THUNK_DATA pThunk = (PIMAGE_THUNK_DATA) 
		   ((PBYTE) hmodCaller + pImportDesc->FirstThunk);

	   // Replace current function address with new function address
	   for (; pThunk->u1.Function; pThunk++) {

		   // Get the address of the function address
		   PROC* ppfn = (PROC*) &pThunk->u1.Function;

		   // Is this the function we're looking for?
		   BOOL fFound = (*ppfn == pfnCurrent);
		   SYSTEM_INFO si;
		   GetSystemInfo(&si);
		   if (!fFound && (*ppfn > si.lpMaximumApplicationAddress)) {

			   // If this is not the function and the address is in a shared DLL, 
			   // then maybe we're running under a debugger on Windows 98. In this 
			   // case, this address points to an instruction that may have the 
			   // correct address.


			   PBYTE pbInFunc = (PBYTE) *ppfn;
			   if (pbInFunc[0] == cPushOpCode) {
				   // We see the PUSH instruction, the real function address follows
				   ppfn = (PROC*) &pbInFunc[1];

				   // Is this the function we're looking for?
				   fFound = (*ppfn == pfnCurrent);
			   }
		   }

		   if (fFound) {
			   // The addresses match, change the import section address
			   WriteProcessMemory(GetCurrentProcess(), ppfn, &pfnNew, 
				   sizeof(pfnNew), NULL);
			   return;  // We did it, get out
		   }
	   }

										   // If we get to here, the function is not in the caller's import section
}